import heapq

from querier.esquerier import ElasticSearchQuerier
from utils import utils
import json


CATEGORIES_TOP_N = 5


class WechatCategoriesQuerier(ElasticSearchQuerier):

    def __init__(self, es, index, doc_type):
        super(WechatCategoriesQuerier, self).__init__(None, None, None)
        self.es = es
        self.index = index
        self.doc_type = doc_type

    def search(self, args):
        id_ = args['biz_code']
        res = self.es.get(index=self.index, doc_type=self.doc_type, id=id_)
        if res['found']:
            categories = res['_source']['categories']
            categories = [utils.category_smzdm_2_decode(int(c)) for c in categories]
            weight = res['_source']['categories_weight']
            len_categories = len(categories)
            len_weight = len(weight)
            s = sum(weight)
            length = len_categories if len_categories <= len_weight else len_weight

            for i in range(0, length):
                if categories[i] in utils.stop_class_0:
                    s = s - weight[i]

            if s <= 0:
                s = 0.1

            ret = [{"text": categories[i], "weight": weight[i]/s} for i in range(0, length) if categories[i] not in utils.stop_class_0]

            return {'categories': ret}
        return {'message': 'biz_code=%s not found.' % id_}

    def _build_query(self, args):
        pass

    def _build_result(self, es_result, param):
        pass
